import cv2
import numpy as np
import time

# 初始化摄像头（0为默认摄像头，可替换为视频文件路径或RTSP地址）
cap = cv2.VideoCapture(0)
if not cap.isOpened():
    raise ValueError("无法打开摄像头")

# 配置摄像头参数（可选）
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)  # 宽度
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)  # 高度
cap.set(cv2.CAP_PROP_FPS, 25)  # 帧率

# 数据增强函数：随机翻转与亮度调整
def data_augmentation(frame, flip_prob=0.5, brightness_scale=0.2):
    # 随机水平翻转
    if np.random.rand() < flip_prob:
        frame = cv2.flip(frame, 1)
    # 调整亮度
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    hsv[..., 2] = np.clip(hsv[..., 2] * (1 + np.random.uniform(-brightness_scale, brightness_scale)), 0, 255)
    return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

# 采集循环（按'q'键退出）
save_path = "./dataset/"  # 数据保存路径
frame_count = 0  # 帧计数器
while True:
    ret, frame = cap.read()
    if not ret:
        break
    
    # 预处理：裁剪中心区域并缩放
    h, w = frame.shape[:2]
    crop_size = min(h, w)
    y = (h - crop_size) // 2
    x = (w - crop_size) // 2
    cropped_frame = frame[y:y+crop_size, x:x+crop_size]
    resized_frame = cv2.resize(cropped_frame, (640, 640))
    
    # 数据增强
    augmented_frame = data_augmentation(resized_frame)
    
    # 可视化与保存（示例：每10帧保存一张图片）
    cv2.imshow("Camera Feed", augmented_frame)
    if frame_count % 10 == 0:
        filename = f"{save_path}/frame_{time.time()}.jpg"
        cv2.imwrite(filename, augmented_frame)
    
    frame_count += 1
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()